import numpy as np
import matplotlib.pyplot as plt
import pandas
import json
import os
import utm
from math import pi, floor
import pickle

# 角度限制, 输入弧度, 输出-pi~pi
def rad_limit(rad: float) -> float:
	if rad >= 180:
		rad = rad - 360
	if rad < -180:
		rad = rad + 360
	return rad

# json to dict?
def json_decode(df_data: pandas.DataFrame):
    columns_name = []
    for key in json.loads(df_data.iloc[0]):
        columns_name.append(key)
    
    temp = pandas.DataFrame(columns=columns_name)
    for i in range(len(df_data)):
        temp_dict = json.loads(df_data.iloc[i])
        temp_series = pandas.Series(temp_dict).to_frame()
        temp = pandas.concat([temp, temp_series.T])
    return temp

def cal_psi_error(data_df: pandas.DataFrame):
    data_df['psi_error'] = data_df['psi'] - data_df['exp_psi']
    for i in range(len(data_df['psi_error'])):
        data_df['psi_error'].iloc[i] = rad_limit(data_df['psi_error'].iloc[i])
    return data_df

if __name__ == '__main__':
    
    HOME = os.getcwd()
    path_x = np.linspace(0, 1000, 200)
    path_y = np.linspace(0, 200, 200)

    exp_line = {'x':path_x,
                'y':path_y}
    exp_line_df = pandas.DataFrame(exp_line)

    with open(f'{HOME}/data_deal/nmpc_path_recorder.pkl', 'rb') as f:
        nmpc_path = pickle.load(f)
        
    with open(f'{HOME}/data_deal/td3_path_recorder.pkl', 'rb') as f:
        td3_path = pickle.load(f)
    
    with open(f'{HOME}/data_deal/sac_path_recorder.pkl', 'rb') as f:
        sac_path = pickle.load(f)  

    mpc_data = pandas.DataFrame(nmpc_path)
    td3_data = pandas.DataFrame(td3_path)
    sacn_data = pandas.DataFrame(sac_path)
    print(sacn_data.columns)
    print(sacn_data)

    plt.rcParams['font.size']=12
    fig, ax = plt.subplots(ncols=3, nrows=4, figsize=(16,14))
    ax[0,0].plot(exp_line_df['x'].to_numpy(), exp_line_df['y'].to_numpy(), color='blue', label='Desire path')
    ax[0,0].plot(mpc_data['path_x'].to_numpy(), mpc_data['path_y'].to_numpy(), color='red', label='NMPC path',alpha=0.7)
    ax[0,0].legend()
    ax[0,0].set_xlabel('x / m')
    ax[0,0].set_ylabel('y / m')

    ax[0,1].plot(exp_line_df['x'].to_numpy(), exp_line_df['y'].to_numpy(), color='blue', label='Desire path')
    ax[0,1].plot(td3_data['path_x'].to_numpy(), td3_data['path_y'].to_numpy(), color='orange',label='TD3 path',alpha=0.7)
    ax[0,1].legend()
    ax[0,1].set_xlabel('x / m')
    ax[0,1].set_ylabel('y / m')

    ax[0,2].plot(exp_line_df['x'].to_numpy(), exp_line_df['y'].to_numpy(), color='blue', label='Desire path')
    ax[0,2].plot(sacn_data['path_x'].to_numpy(), sacn_data['path_y'].to_numpy(), color='green', label='SAC-N path',alpha=0.7)
    ax[0,2].legend()
    ax[0,2].set_xlabel('x / m')
    ax[0,2].set_ylabel('y / m')

    ax[1,0].plot(mpc_data['y_error'].to_numpy(), color='red', label='NMPC')
    ax[1,0].legend()
    ax[1,0].set_ylabel('Cross-track error / m')
    ax[1,0].set_xlabel('Time / s')

    ax[1,1].plot(td3_data['y_error'].to_numpy(), color='orange', label='TD3')
    ax[1,1].legend()
    ax[1,1].set_ylabel('Cross-track error / m')
    ax[1,1].set_xlabel('Time / s')

    ax[1,2].plot(sacn_data['y_error'].to_numpy(), color='green', label='SAC-N')
    ax[1,2].legend()
    ax[1,2].set_ylabel('Cross-track error / m')
    ax[1,2].set_xlabel('Time / s')

    
    ax[2,0].plot(mpc_data['exp_anle'].to_numpy() - mpc_data['yaw'].to_numpy() , color='red', label='NMPC')
    ax[2,0].legend()
    ax[2,0].set_xlabel('Time / s')
    ax[2,0].set_ylabel('Heading angle error / $^o$')
    ax[2,0].set_ylim(-12,52)
    ax[2,1].plot(td3_data['exp_anle'].to_numpy() - td3_data['yaw'].to_numpy() , color='orange', label='TD3')
    ax[2,1].legend()
    ax[2,1].set_xlabel('Time / s')
    ax[2,1].set_ylabel('Heading angle error / $^o$')
    ax[2,1].set_ylim(-12,52)
    ax[2,2].plot(sacn_data['exp_anle'].to_numpy() - sacn_data['yaw'].to_numpy() , color='green', label='SAC-N')
    ax[2,2].legend()
    ax[2,2].set_xlabel('Time / s')
    ax[2,2].set_ylabel('Heading angle error / $^o$')
    ax[2,2].set_ylim(-12,52)
    
    ax[3,0].plot(mpc_data['rudder'].to_numpy(), color='red', label='NMPC')
    ax[3,0].legend()
    ax[3,0].set_ylabel('Steering command / $^o$')
    ax[3,0].set_xlabel('Time / s')
    ax[3,0].set_ylim(-15,32)

    ax[3,1].plot(td3_data['rudder'].to_numpy(), color='orange', label='TD3')
    ax[3,1].legend()
    ax[3,1].set_ylabel('Steering command / $^o$')
    ax[3,1].set_xlabel('Time / s')
    ax[3,1].set_ylim(-15,32)

    ax[3,2].plot(sacn_data['rudder'].to_numpy(), color='green', label='SAC-N')
    ax[3,2].legend()
    ax[3,2].set_ylabel('Steering command / $^o$')
    ax[3,2].set_xlabel('Time / s')
    ax[3,2].set_ylim(-15,32)


    print(np.mean(np.abs(mpc_data['y_error'].to_numpy())))
    print(np.mean(np.abs(td3_data['y_error'].to_numpy())))
    print(np.mean(np.abs(sacn_data['y_error'].to_numpy())))
    
    print(np.mean(np.abs(mpc_data['exp_anle'].to_numpy() - mpc_data['yaw'].to_numpy())))
    print(np.mean(np.abs(td3_data['exp_anle'].to_numpy() - td3_data['yaw'].to_numpy())))
    print(np.mean(np.abs(sacn_data['exp_anle'].to_numpy() - sacn_data['yaw'].to_numpy())))

    print(np.mean(np.abs(mpc_data['rudder'].to_numpy())))
    print(np.mean(np.abs(td3_data['rudder'].to_numpy())))
    print(np.mean(np.abs(sacn_data['rudder'].to_numpy())))

    plt.savefig(f'{HOME}/data_deal/plot_path_simulation.eps', bbox_inches='tight', dpi=300)
    plt.show()